package com.csw.android.dev_utils.ex

import androidx.lifecycle.ViewModel
import androidx.lifecycle.ViewModelProvider
import androidx.lifecycle.ViewModelStoreOwner
import java.lang.ref.WeakReference
import java.util.LinkedList

/** viewModelStore创建的ViewModel的记录 */
private class ViewModelStoreRecord(
    viewModelStoreOwner: ViewModelStoreOwner,
) {
    val weakViewModelStoreOwner = WeakReference(viewModelStoreOwner)
    val weakViewModels = HashMap<Class<*>, WeakReference<ViewModel>>()
}

/** 所有viewModel的示例持有对象 */
private val instanceHolder = LinkedList<ViewModelStoreRecord>()

/**
 * 查询ViewModelStoreRecord，如果不存在则创建一个，遍历过程中，对已经回收的记录对象进行移除
 */
@Synchronized
private fun ensureViewModelStoreRecord(viewModelStoreOwner: ViewModelStoreOwner): ViewModelStoreRecord {
    instanceHolder.iterator().run {
        while (hasNext()) {
            next().run {
                if (weakViewModelStoreOwner.get() === viewModelStoreOwner) {
                    return this
                } else if (weakViewModelStoreOwner.get() == null) {
                    remove()
                }
            }
        }
    }
    return ViewModelStoreRecord(viewModelStoreOwner).apply {
        instanceHolder.add(this)
    }
}

/**
 * 取得Store中的ViewModel，若不存在，则创建
 */
fun <T : ViewModel> ViewModelStoreOwner.get(clazz: Class<T>): T {
    ViewModelProvider(this)[clazz].let { viewModel ->
        ensureViewModelStoreRecord(this).weakViewModels[clazz] = WeakReference(viewModel)
        return viewModel
    }
}

/**
 * 查询Store中的viewModel，不存在返回空
 */
fun <T : ViewModel> ViewModelStoreOwner.query(clazz: Class<T>): T? {
    return try {
        ensureViewModelStoreRecord(this).weakViewModels[clazz]?.get() as? T
    } catch (e: Exception) {
        e.printStackTrace()
        null
    }
}